import cv2 as cv
import numpy as np
import time

def make_square(img):
    x, y = img.shape
    size = max(x, y)
    new_img = np.zeros((size, size), dtype=np.uint8)
    ax,ay = (size - img.shape[1])//2,(size - img.shape[0])//2
    new_img[ay:img.shape[0]+ay,ax:ax+img.shape[1]] = img
    # Pad the new_img array with edge pixel values
    # Apply feathering effect
    feather_pixels = 8
    for i in range(feather_pixels):
        alpha = (i + 1) / feather_pixels
        new_img[ay - i - 1, :] = img[0, :] * (1 - alpha)  # Top edge
        new_img[ay + img.shape[0] + i, :] = img[-1, :] * (1 - alpha)  # Bottom edge
    return new_img

# 找到发光点坐标
def _find_dot(img):
    # img = cv.GaussianBlur(img,(5,5),0)
    # grey_img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
    # 图像二值化 [1]表示只接受元组索引为1的返回值
    # binary_img = cv.threshold(grey_img, 255*0.8, 255, cv.THRESH_BINARY)[1]
    # 找出轮廓
    # print(img.format)
    # 发光点轮廓检测
    contours = cv.findContours(img, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)[1]
    # 画出轮廓
    img = cv.drawContours(img, contours, -1, 255, 1) # (0,255,0)

    image_points = []
    for contour in contours:
        moments = cv.moments(contour)
        if moments["m00"] != 0:
            center_x = int(moments["m10"] / moments["m00"])
            center_y = int(moments["m01"] / moments["m00"])
            cv.putText(img, f'({center_x}, {center_y})', (center_x,center_y - 15), cv.FONT_HERSHEY_SIMPLEX, 0.3, (100,255,100), 1)
            cv.circle(img, (center_x,center_y), 1, (100,255,100), -1)
            image_points.append([center_x, center_y])

    if len(image_points) == 0:
        image_points = [[None, None]]

    return img, image_points

print(cv.__version__)
last_time = 0
now_time = 0
cap = cv.VideoCapture('output1.avi')
width = 1280  # 定义摄像头获取图像宽度
height = 960  # 定义摄像头获取图像长度
fps = 60
cap.set(6, cv.VideoWriter_fourcc('M', 'J', 'P', 'G')) # 设置图像的编码格式
cap.set(cv.CAP_PROP_FPS, fps)  # 帧率
cap.set(cv.CAP_PROP_FRAME_WIDTH, width)  # 设置宽度
cap.set(cv.CAP_PROP_FRAME_HEIGHT, height)  # 设置高度

#设置图像保存格式与路径
# fourcc = cv.VideoWriter_fourcc('M','J','P','G')
# out = cv.VideoWriter('output.avi',fourcc, 60.0, (1920,1080))
i = 0
kernel = np.array([[-2,-1,-1,-1,-2],
                        [-1,1,3,1,-1],
                        [-1,3,4,3,-1],
                        [-1,1,3,1,-1],
                        [-2,-1,-1,-1,-2]])
# 定义HSV阈值
hsv_low = np.array([  0, 0,149])
hsv_high= np.array([180,39,255])

while(cap.isOpened()):
# while True:
    last_time = now_time
    now_time = time.time()
    p = now_time - last_time
    fps = 1 / p

    ret, frame = cap.read()
    if not ret:
        print("Can't receive frame (stream end?). Exiting ...")
        break

    if cv.waitKey(1) == ord('s'):
        cv.imwrite('./pic/pic'+str(i)+'.jpg', frame)
        i = i+1
    
    # frame = cv.GaussianBlur(frame,(9,9),0)
    # frame = cv.filter2D(frame, -1, kernel)
    # frame = cv.cvtColor(frame, cv.COLOR_RGB2BGR)
    
    frame = cv.cvtColor(frame, cv.COLOR_RGB2HSV) # 转为HSV色彩空间
    mask = cv.inRange(frame, hsv_low, hsv_high)  # 通过HSV的高低阈值，提取图像部分区域，输出的mask为二值图像
    mask = make_square(mask)
    # 调用_find_dot画出轮廓和每个点坐标
    mask, single_camera_image_points = _find_dot(mask)
    # frame = cv.cvtColor(frame, cv.COLOR_RGB2GRAY)
    # frame = cv.GaussianBlur(frame,(9,9),0)
    
    # 在图像上打印帧率
    cv.putText(mask, "{:.2f}".format(fps), (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv.LINE_AA)
    
    
    cv.imshow('frame', mask)

    if cv.waitKey(1) == ord('q'):
        break
    
    
# Release everything if job is finished
cap.release()
# out.release()
cv.destroyAllWindows()
